# 多项式回归 用 statsmodels  6次多项式的数据
import numpy as np
import matplotlib.pyplot as plt
import statsmodels.api as sm
from sklearn.metrics import mean_squared_error

x = np.linspace(-1, 1, 101).reshape(-1, 1)
num_cofffs = 6
coeffs = [1, 2, 3, 4, 5, 6]
Y_true = 0
for i in range(num_cofffs):
    Y_true += coeffs[i] * np.power(x, i)

print(Y_true.shape)
print(x.shape)

Y = Y_true + np.random.randn(*x.shape) * 1.5
plt.figure(1)
plt.scatter(x, Y)
# plt.show()
# plt.pause(1)


x2 = x ** 2
x3 = x ** 3
x4 = x ** 4
x5 = x ** 5
XX = np.hstack((x, x2, x3, x4, x5))

XX = sm.add_constant(XX)
model = sm.OLS(Y, XX).fit()

pred = model.fittedvalues
rmse = np.sqrt(mean_squared_error(Y, pred))
print(rmse)
print(model.summary())
w = model.params

plt.figure(1)
plt.plot(x, pred, 'r', label='predict')
plt.plot(x, Y_true, 'b', label='true')
plt.legend()
plt.show()
